from __future__ import print_function
from __future__ import division
import sys
import os
import csv
import time
import argparse

import serial

import sr245
import zpickle

# Preferences
# -----------
RAW_DATA_EXT = os.path.extsep.join(['pkl','gz'])
OUTPUT_FILE_EXT = 'csv'

# Experiment Setup
# ----------------
# approximate capture rate of setup, used for break detection calculation 
TRIG_RATE = 1300 # Hz
assert TRIG_RATE < sr245.MAX_RATE
SWEEP_RATE = 2 # Hz
TRIG_ENABLE_PIN = 2
TRIG_ENABLE_POL = 0 # hardware level that *enables* triggers

# Default Capture Settings
# ----------------
SWEEPS = 20
MIN_CAP_LEN = 250 # average seems to be about 271-274 samples

# Computer Settings
# -----------------
PORT = 'COM3'
BAUD = 19200
SERIAL_FORMAT = '8N1'
# Break detection threshold (space between scans)
BREAK_THOLD_MIN = 0.05 # seconds.
BREAK_THOLD = max(3 * max(1/TRIG_RATE, 10/BAUD), BREAK_THOLD_MIN)
# - minimum is from BREAK_THOLD_MIN var
# - otherwise at least 3 (arbitrary fudge factor) times
#    - byte rate (10 symbols/byte)
#    - approximate trigger rate (not measured, just provided)

# DEBUGGING
# ---------
DBG_VERBOSE = 0
DBG_USESAVED = False
DBG_SAVEDATA = False
DBG_PLOT = False

if DBG_PLOT: 
    import pylab
    
class ScanSettings(object):
    def __init__(self):
        #self.file
        #self.port
        #self.baudrate
        #self.serial_format
        #self.sweeps
        #self.save_binary
        #self.to_stdio
        pass
            
def main(argv=None):
    if argv is None:
        argv = sys.argv
        
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter, 
        description='''
Captures data from the SR245 and saves it to a CSV file for use by GENPLOT.  

If the -f argument is not provided, the program will run in an interactive 
mode, asking for an output file name every iteration.  To configure the 
settings in the interactive mode (number of sweeps, serial port), use the 
other switches when launching the program.

Interactive mode can be quit with entering a filename of "q", or using Ctrl-C
to kill the program.  After capturing data to a provided file name, it will 
increment the last character (be it number or letter) and offer that as a guess
for what to save the next sweep to.  Hitting Return/Enter, leaving the prompt
blank, will accept that and run another sweep, or any other name can be entered
to override the guess.
'''.strip())
    # positional argument converted to optional: if it's not provided, that 
    # means run in interactive mode
    #parser.add_argument('file',
    #    help='output file name.  Extension defaults to ".csv" if not provided.')
    parser.add_argument('-f', dest='file', metavar='FILENAME', type=str,
        help=('output file name, extension defaults to ".{}" if not provided.')
            .format(OUTPUT_FILE_EXT))
    parser.add_argument('-s', dest='port', metavar='PORT', type=str,
        help=('serial port where the data acquisition module (DAQ) is '
              'connected, default: {:s}')
            .format(PORT), default=PORT)
    parser.add_argument('-sb', dest='baudrate', metavar='BAUDRATE', type=int,
        help='serial baud rate, default: {:d}'
            .format(BAUD), default=BAUD)
    parser.add_argument('-sf', dest='serial_format', metavar='BYTE_FORMAT', type=str,
        help='serial format, default: {:s}'
            .format(SERIAL_FORMAT), default=SERIAL_FORMAT)
    parser.add_argument('-n', dest='sweeps', metavar='SWEEPS', type=int,
        help='number of sweeps to capture, default: {:d}'
            .format(SWEEPS), default=SWEEPS)
    parser.add_argument('-p', dest='plot', action='store_true', 
        help='display plot of data')
    parser.add_argument('-r', dest='save_binary', action='store_false', 
        help='prevents saving raw data')
    #parser.add_argument('-s', dest='to_stdio', action='store_true', 
    #    help=('NOT IMPLEMENTED -- '
    #          'output is redirected to standard I/O (printed to console), '
    #          'and no formatted output file is saved.'))
    
    settings = parser.parse_args(argv[1:])
        
    if settings.file is not None:    
        # if "-f filename" provided, run once with provided settings
        capture_data(settings)
    else:
        # run interactively
        print('INTERACTIVE MODE\n----------------')
        
        while True:
            if settings.file is None:
                settings.file = raw_input('Enter output file name:  ')
            else:
                # Guess the next file by incrementing last char of previous
                nextfile = settings.file[:-1] + chr(ord(settings.file[-1]) + 1)
                settings.file = raw_input(
                    'Filename to run again [%s]: ' % nextfile)
                # if nothing entered, assume guess was correct
                settings.file = settings.file or nextfile
            
            # exit with 'q', 'Q' or Ctrl-C
            if settings.file.upper() == 'Q': 
                print ('\nExit.')
                return
            
            capture_data(settings)
        
def capture_data(settings):

    # do the annoyingly long import if requested
    if settings.plot: import pylab
        
    # generate output file names
    parts = settings.file.split(os.path.extsep)
    if len(parts) > 1 and (2 <= len(parts[-1]) <= 4):
        # filename provided has extension
        out_file_csv = settings.file
        # rip off extension to provide our own
        out_file_raw = os.path.extsep.join(parts[:-1] + [RAW_DATA_EXT])
    else:
        # no extension, all up to us
        out_file_csv = settings.file + os.path.extsep + OUTPUT_FILE_EXT
        out_file_raw = settings.file + os.path.extsep + RAW_DATA_EXT
    
    if not DBG_USESAVED:
        #q = sr245.sr245(serial.Serial(PORT, BAUD))
        
        q = sr245.sr245(
                serial.Serial(
                    port=settings.port, 
                    baudrate=settings.baudrate,
                    bytesize=int(settings.serial_format[0]),
                    parity=settings.serial_format[1].upper(),
                    stopbits=float(settings.serial_format[2:])
                )
            )
        # trigger enable/disable aliases
        en_trig = lambda: q.write_d(TRIG_ENABLE_PIN, TRIG_ENABLE_POL)
        dis_trig = lambda: q.write_d(TRIG_ENABLE_PIN, not TRIG_ENABLE_POL)
        
        dis_trig()
        q.sync_mode()
        
        # allow D flip-flop time to shut up triggers--should take 0.5 s max
        time.sleep(1/SWEEP_RATE * 1.5)
        
        print('Sweep run: %s' % out_file_csv)
            
        iters = []
        for i in range(settings.sweeps):
            attempt = 0
            while True:
                print(' - sweep {:3d} of {:3d}'
                    .format(i+1, settings.sweeps), end='')
                if attempt:
                    print(', attempt {}'
                        .format(attempt+1), end='')
                
                en_trig()
                q.unmask_trigger()
                # the length of the scan is the maximum number of samples, but 
                # the timeout will likely make it end up being shorter.  the 
                # actual number it usually takes is about 275, so 300+ would 
                # be good.
                cap = q.scan(500, stream=True, timeout=BREAK_THOLD, 
                                cb_running=dis_trig)
                q.mask_trigger()
                
                iters.append(cap)
                
                if len(cap) < MIN_CAP_LEN:
                    print(' - ERROR: too short (<{})'.format(MIN_CAP_LEN))
                    # discard later, but save for raw data file
                else:
                    print('\r'+' '*70,end='\r') # print over line to clear
                    break
                    
                attempt += 1
                
        q.scan_end()
        dis_trig()
        
        q.disconnect()

        if DBG_SAVEDATA:
            #with open('debug-dump.bin','wb') as f:
            #    pickle.dump(iters, f)
            zpickle.save(iters, 'debug-dump.bin')
                            
        if settings.save_binary:
            #with open(out_file + '.' + RAW_DATA_EXT,'wb') as f:
            #    pickle.dump(iters, f)
            zpickle.save(iters, out_file_raw)
                
    else:
        #with open('debug-dump.bin','rb') as f:
        #    iters = pickle.load(f)
        iters = zpickle.load('debug-dump.bin')
            
    if settings.plot or DBG_PLOT:
        ## stack all the sweeps on each other and display (average)
        accum = [sum(x) for x in zip(*iters)]
        pylab.plot(range(len(accum)), accum)
        pylab.show()
        
    # drop too-short sweeps
    iters = filter(lambda x: len(x) >= MIN_CAP_LEN, iters)
    
    # output csv data
    # - create info header
    info = []
    info.append('SR245 Scan Data')
    info.append('Captured on: {}'.format(
        time.strftime("%a, %d %b %Y %H:%M:%S +0000", time.gmtime())))
    info.append('Sweeps (requested): {}'.format(settings.sweeps))
    info.append('Sweeps (used)     : {}'.format(len(iters)))
    info.append('Raw data          : {}'.format(
        ('saved to ' + out_file_raw) if settings.save_binary else 'not saved')
    )
    info.append('')
    info.append('Fields:')
    info.append('Index, Sweep 1, Sweep 2, ..., Sweep n')
    
    # - append info lines with 'C', commenting them out in GENPLOT
    info = ['C\t{}{}'.format(line, os.linesep) for line in info]
    
    # - trim data sets so they are all the same length
    trimmed = zip(*iters)
    indexcol = range(1,1+len(trimmed))
    dataout = zip(indexcol, *zip(*trimmed))
    
    # - write file
    with open(out_file_csv, 'wb') as f:
        f.writelines(info)
        w = csv.writer(f)
        w.writerows(dataout)
            
    print(' - complete, saved to %s' % out_file_csv)
    return

if __name__ == '__main__':
    sys.exit(main())